#!/usr/bin/env python3
import numpy as np
from typing import Dict, Any
import sys
import os
from pathlib import Path
# Fix for Python 3.10 compatibility - ensure path is string
robottwin_path = str(Path(__file__).parent.parent / 'RoboTwin')
if robottwin_path not in sys.path:
    sys.path.append(robottwin_path)
from envs.utils.vis_pcd import PointArrow, catpcd
from scalePcd import scalePcd
from keyPcd import keypose_item
from scipy.spatial.transform import Rotation

def trajPcd(unifiedData: Dict[str, Any], param: str = 's', *, surpress_print: bool = False) -> Dict[str, Any]:
    """
    Convert trajectory data (state/action) to point cloud arrows and inject into existing pointclouds.
    about left and right: the action / state 's 0:7 is the left on first view, and keypose's left is also left on first view
    but the semantic part's left is third view's left, so here the current_bias l and r are second view's and when adding to state/action, it needs to be crossed.
    """
    assert param in ['s', 'a','e','m','i'], f"param must be 's' (state) or 'a' (action) or 'e' (modifying the state and action inplace to use two finger end's middle point), got {param}"
    if param in ["s","a"]:
        return add_js_ja(unifiedData, param, surpress_print)
    elif param == "e":
        return add_mid_end(unifiedData, param, surpress_print)
    elif param == "m":
        return move_state(unifiedData, param, surpress_print)
    elif param == "i":
        return move_state(unifiedData, param, surpress_print, inv=True)
    print(f"    ❌ TRAJPCD({param}): Invalid parameter, got {param}")
    return unifiedData 
def move_state_7D(state_7D, local_vec, inv=False):
    """Move 7D pose by local vector in local frame, return modified global pose"""
    pos, quat = state_7D[:3], state_7D[3:]
    local_vec=np.array(local_vec)
    rotation = Rotation.from_quat(quat)
    rotated_vec = rotation.apply(-local_vec if inv else local_vec)
    return np.concatenate([pos + rotated_vec, quat])
def move_state(unifiedData: Dict[str, Any], param: str = 'm',  surpress_print: bool = False, inv=False) -> Dict[str, Any]:
    """Move both arms' states by local vectors"""
    print("original state inv",inv,list(np.round(unifiedData["state"][:3],3)),list(np.round(unifiedData["state"][8:11],3)))
    unifiedData["state"][:7] = move_state_7D(unifiedData["state"][:7], [-0.12,0,0], inv)
    unifiedData["state"][8:15] = move_state_7D(unifiedData["state"][8:15], [-0.12,0,0], inv)
    print("moved state inv",inv,list(np.round(unifiedData["state"][:3],3)),list(np.round(unifiedData["state"][8:11],3)))
    return unifiedData
def add_mid_end(unifiedData: Dict[str, Any], param: str = 'e',  surpress_print: bool = False) -> Dict[str, Any]:
    assert "key_poses" in unifiedData and unifiedData["key_poses"] is not None, f"No key_poses data available for trajPcd"
    assert "state" in unifiedData and unifiedData["state"] is not None, f"No state data available for trajPcd"
    assert "action" in unifiedData and unifiedData["action"] is not None, f"No action data available for trajPcd"
    keypose_data = unifiedData['key_poses']
    llend=keypose_item(keypose_data)["left_gripper"]["left_finger_end"]
    lrend=keypose_item(keypose_data)["left_gripper"]["right_finger_end"]
    rlend=keypose_item(keypose_data)["right_gripper"]["left_finger_end"]
    rrend=keypose_item(keypose_data)["right_gripper"]["right_finger_end"]
    lend_mid_xyz=((np.array(lrend)+np.array(llend))/2)
    lend_mid_rpy=keypose_item(keypose_data)["left_gripper"]["rpy"]
    rend_mid_xyz=((np.array(rlend)+np.array(rrend))/2)
    rend_mid_rpy=keypose_item(keypose_data)["right_gripper"]["rpy"]
    lend_midpose=np.concatenate([lend_mid_xyz,lend_mid_rpy])
    rend_midpose=np.concatenate([rend_mid_xyz,rend_mid_rpy])
    mid_end_pcd_1=PointArrow(lend_midpose, [0, 1, 0], 100)  # green
    mid_end_pcd_2=PointArrow(rend_midpose, [0, 0, 1], 100)  # blue
    mid_end_pcd=catpcd(mid_end_pcd_1, mid_end_pcd_2)
    # print("l_kp_pos-pose_data_1[:3]",np.linalg.norm(l_kp_pos-pose_data_1[:3]))
    # print("r_kp_pos-pose_data_2[:3]",np.linalg.norm(r_kp_pos-pose_data_2[:3]))
    # print("middlePoint",((np.array(llend)+np.array(rrend))/2))
    # print("left hand middlePoint-state distance",np.linalg.norm((np.array(lrend)+np.array(llend))/2-pose_data_1[:3]))
    # print("right hand middlePoint-state distance",np.linalg.norm((np.array(rrend)+np.array(rlend))/2-pose_data_2[:3]))
    return plot_arrow(unifiedData, mid_end_pcd, param, surpress_print)
def add_js_ja(unifiedData: Dict[str, Any], param: str ,surpress_print: bool = False) -> Dict[str, Any]:
    data_key = {'s': 'state', 'a': 'action'}[param]
    assert data_key in unifiedData and unifiedData[data_key] is not None, f"No {data_key} data available for trajPcd"
    trajectory_data = unifiedData[data_key]
    assert len(trajectory_data.shape) == 1, f"Expected 1D trajectory data, got shape {trajectory_data.shape}"
    assert trajectory_data.shape[0]==16,f"Expected 16D trajectory data, got shape {trajectory_data.shape}"
    pose_data_1 = trajectory_data[:7]  # Take first 7 elements as pose
    pose_data_2 = trajectory_data[8:15]  # Take first 7 elements as pose
    traj_pcd_1 = PointArrow(pose_data_1, [1, 0, 1] if param == 's' else [0, 1, 1], 100)  # Magenta for state, Cyan for action on the 0:7 side of 16 D
    traj_pcd_2 = PointArrow(pose_data_2, [0.5, 0, 0.5] if param == 's' else [0, 0.5, 0.5], 100)  # Dark red = for state,Dark blue for action on the 8:15 side of 16 D
    traj_pcd = catpcd(traj_pcd_1, traj_pcd_2)
    return plot_arrow(unifiedData, traj_pcd, param, surpress_print)
def plot_arrow(unifiedData: Dict[str, Any], traj_pcd, param: str,surpress_print: bool = False):
    traj_pcd = scalePcd(traj_pcd, tag=f"traj_{param}")
    updated_fields = {}
    for pcd_key in unifiedData.keys():
        if "pointcloud" in pcd_key and unifiedData[pcd_key] is not None:
            existing_pcd = unifiedData[pcd_key]
            assert len(existing_pcd.shape) == 2 and existing_pcd.shape[1] == 6, f"Skipping {pcd_key}: Expected shape (Np, 6), got {existing_pcd.shape}"
            original_count = existing_pcd.shape[0]
            existing_pcd = scalePcd(existing_pcd, tag="existing_pcd")
            traj_pcd_to_use = scalePcd(traj_pcd, tag="traj_pcd_to_use")
            combined_pcd = catpcd(existing_pcd, traj_pcd_to_use)
            unifiedData[pcd_key] = combined_pcd
            updated_fields[pcd_key] = (original_count, combined_pcd.shape[0])
    
    for key, (original_count, final_count) in updated_fields.items():
        if not surpress_print:
            print(f"    🎯 TRAJPCD({param},{key}): {original_count} → {final_count}")
    return unifiedData